// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <functional>
#include <random>

#include <melon/random.h>
#include <melon/range.h>

namespace kumo::pollux::functions::kll {
    namespace detail {
        /// Internal API, do not use outside Pollux.
        template<typename T>
        struct View {
            uint32_t k;
            size_t n;
            T minValue;
            T maxValue;
            melon::Range<const T *> items;
            melon::Range<const uint32_t *> levels;

            uint8_t numLevels() const {
                return levels.size() - 1;
            }

            uint32_t safeLevelSize(uint8_t level) const {
                return level < numLevels() ? levels[level + 1] - levels[level] : 0;
            }

            void deserialize(const char * MELON_NONNULL);
        };
    } // namespace detail

    constexpr uint32_t kDefaultK = 200;

    /// Estimate the proper k value to ensure the error bound epsilon.
    uint32_t kFromEpsilon(double epsilon);

    /// Implementation of KLL sketch that can achieve nearly optimal
    /// accuracy per retained item.
    ///
    /// This sketch is configured with a parameter k, which affects the
    /// size of the sketch and its estimation error.
    ///
    /// The estimation error is commonly called epsilon and is a fraction
    /// between zero and one.  Larger values of k result in smaller values
    /// of epsilon.  Epsilon is always with respect to the quantile and
    /// cannot be applied to the corresponding values.
    ///
    /// The default k of 200 yields a "single-sided" epsilon of about
    /// 1.33%.
    ///
    /// See https://arxiv.org/abs/1603.05346v2 for more details.
    template<
        typename T,
        typename Allocator = std::allocator<T>,
        typename Compare = std::less<T> >
    struct KllSketch {
        KllSketch(
            uint32_t k = kll::kDefaultK,
            const Allocator & = Allocator(),
            uint32_t seed = melon::Random::rand32());

        /// Cannot be called after insert().
        void setK(uint32_t k);

        /// Add one new value to the sketch.
        void insert(T value);

        /// Call this before serialization can optimize the space used.
        void compact();

        /// Merge this sketch with values from multiple other sketches.
        /// @tparam Iter Iterator type dereferenceable to the same type as this sketch
        ///  (KllSketch<T, Allocator, Compare>)
        /// @param sketches Range of sketches to be merged to this one
        template<typename Iter>
        void merge(const melon::Range<Iter> &sketches);

        /// Merge with another sketch.
        void merge(const KllSketch<T, Allocator, Compare> &other);

        /// Must be called before estimateQuantile() or estimateQuantiles().
        void finish();

        /// Estimate the value of the given quantile.
        /// @param quantile Quantile in [0, 1] to be estimated
        T estimateQuantile(double quantile) const;

        /// Estimate the values of the given quantiles.  This is more
        /// efficient than calling estimateQuantile(double) repeatedly.
        /// @tparam Iter Iterator type dereferenceable to double
        /// @param quantiles Range of quantiles in [0, 1] to be estimated
        template<typename Iter>
        std::vector<T, Allocator> estimateQuantiles(
            const melon::Range<Iter> &quantiles) const;

        /// Estimate the values of the given quantiles.  This is more
        /// efficient than calling estimateQuantile(double) repeatedly.
        /// @tparam Iter Iterator type dereferenceable to double
        /// @param quantiles Range of quantiles in [0, 1] to be estimated
        /// @param out Pre-allocated memory to hold the result, must be at least as
        ///  large as `quantiles`
        template<typename Iter>
        void estimateQuantiles(const melon::Range<Iter> &quantiles, T *out) const;

        /// The total number of values being added to the sketch.
        size_t totalCount() const {
            return n_;
        }

        uint32_t k() const {
            return k_;
        }

        /// Calculate the size needed for serialization.
        size_t serializedByteSize() const;

        /// Serialize the sketch into bytes.  The serialzation is versioned, and newer
        /// version of code should be able to read all previous versions.
        ///
        /// @param out Pre-allocated memory at least serializedByteSize() in size
        void serialize(char *out) const;

        /// Deserialize a sketch from bytes.  Newer version of code should be able to
        /// read all previous versions.
        static KllSketch<T, Allocator, Compare> deserialize(
            const char *data,
            const Allocator & = Allocator(),
            uint32_t seed = melon::Random::rand32());

        /// Equivalent to calling insert(value) `count` times on a new
        /// sketch, but more efficient.
        static KllSketch<T, Allocator, Compare> fromRepeatedValue(
            T value,
            size_t count,
            uint32_t k = kll::kDefaultK,
            const Allocator & = Allocator(),
            uint32_t seed = melon::Random::rand32());

        /// Merge this sketch with values from multiple other deserialized sketches.
        /// @tparam Iter Iterator type dereferenceable to `const char*`, which
        ///  represents a deserialized sketch
        /// @param sketches Range of sketches to be merged to this one
        template<typename Iter>
        void mergeDeserialized(const melon::Range<Iter> &sketches);

        /// Merge with another deserialized sketch.  This is more efficient
        /// than deserialize then merge.
        void mergeDeserialized(const char *data);

        /// Get frequencies of items being tracked.  The result is sorted by item.
        std::vector<std::pair<T, uint64_t> > getFrequencies() const;

        /// Internal API, do not use outside Pollux.
        void mergeViews(const melon::Range<const detail::View<T> *> &views);

        /// Internal API, do not use outside Pollux.
        detail::View<T> toView() const;

        static KllSketch<T, Allocator, Compare>
        fromView(const detail::View<T> &, const Allocator &allocator, uint32_t seed);

    private:
        KllSketch(const Allocator &, uint32_t seed);

        void doInsert(T);

        uint32_t insertPosition();

        int findLevelToCompact() const;

        void addEmptyTopLevelToCompletelyFullSketch();

        void shiftItems(uint32_t delta);

        uint8_t numLevels() const {
            return levels_.size() - 1;
        }

        uint32_t getNumRetained() const {
            return levels_.back() - levels_[0];
        }

        uint32_t safeLevelSize(uint8_t level) const {
            return level < numLevels() ? levels_[level + 1] - levels_[level] : 0;
        }

        using AllocU32 = typename std::allocator_traits<
            Allocator>::template rebind_alloc<uint32_t>;

        uint32_t k_;
        Allocator allocator_;

        // mt19937 uses too much memory (up to 5000 bytes), we choose to use
        // default_random_engine here to sacrifice some randomness for memory.
        std::independent_bits_engine<std::default_random_engine, 1, uint32_t>
        randomBit_;

        size_t n_;
        T minValue_;
        T maxValue_;
        std::vector<T, Allocator> items_;
        std::vector<uint32_t, AllocU32> levels_;
        bool isLevelZeroSorted_;
    };
} // namespace kumo::pollux::functions::kll

#include <pollux/functions/lib/kll_sketch-inl.h>
